﻿using System;
using System.Threading;
using System.Threading.Tasks;

namespace OmniKits.Core
{
    using Extensions;

    public abstract class SimpleRunTaskBase<TTask> : IRunTask<TTask>
        where TTask : Task
    {
        private object _SyncToken = new object();

        // necessary to be volatile for double-checked locking
        private volatile TTask _Task = null;
        public TTask CurrentTask => _Task;

        protected bool IsNotRunning => CurrentTask.IsNotRunning();
        public bool IsRunning => CurrentTask.IsRunning();

        protected abstract TTask StartTask(CancellationToken cancellationToken);

        public TTask Run(CancellationToken cancellationToken)
        {
            var task = CurrentTask;
            if (task.IsNotRunning()) // first check
            {
                lock (_SyncToken) // locked
                {
                    task = CurrentTask;
                    if (task.IsNotRunning())  // second check
                    {
                        task = StartTask(cancellationToken);
                        _Task = task;
                        return task;
                    }
                }
            }

            if (cancellationToken.CanBeCanceled)
                throw new InvalidOperationException();

            return task;
        }
    }

    public abstract class CanStopRunTaskBase<TTask> : IRunTask<TTask>, ICanStop
        where TTask : Task
    {
        private class TaskState
        {
            public CancellationTokenSource CancellationTokenSource { get; }
            public TTask Task { get; }

            public TaskState(TTask task, CancellationTokenSource cts)
            {
                CancellationTokenSource = cts;
                Task = task;
            }

            public bool Register(CancellationToken token)
            {
                if (token.CanBeCanceled)
                {
                    token.Register(CancellationTokenSource.Cancel);
                    return true;
                }
                return false;
            }
        }

        private volatile TaskState _State;
        private object _SyncToken = new object();

        public TTask CurrentTask => _State?.Task;

        protected bool IsNotRunning => CurrentTask.IsNotRunning();
        public bool IsRunning => CurrentTask.IsRunning();

        protected abstract TTask StartTask(CancellationToken cancellationToken);

        public TTask Run(CancellationToken cancellationToken)
        {
            TaskState state = _State;
            if ((state?.Task).IsNotRunning())
            {
                lock (_SyncToken)
                {
                    state = _State;
                    if ((state?.Task).IsNotRunning())
                    {
                        var cts = new CancellationTokenSource();
                        if (cancellationToken.CanBeCanceled)
                            cancellationToken.Register(cts.Cancel);

                        var task = StartTask(cts.Token);
                        state = new TaskState(task, cts);
                        _State = state;
                        return state.Task;
                    }
                }
            }

            state.Register(cancellationToken);
            return state.Task;
        }

        public bool Stop()
        {
            var state = Interlocked.Exchange(ref _State, null);

            if (state == null) return false;

            var done = state.Task.IsCompleted;

            state.CancellationTokenSource.Cancel();

            return !done;
        }
    }
}
